from io import BytesIO

from fastapi import UploadFile
from paddleocr import PaddleOCR

from output.tools import get_data_text
from tbpu import get_parser
from PIL import Image
import numpy as np

from server import FileParser, DEFAULT_TBPU

paddle = PaddleOCR(
    det_limit_side_len=1200,
    use_gpu=True,
    gpu_id=0
)

class ImgParser(FileParser):

    extensions = ['jpg', 'jpeg', 'png', 'gif', 'tiff', 'bmp', 'webp']

    async def parse_upload_file_to_text(self, file: UploadFile, mode=DEFAULT_TBPU):
        contents = await file.read()
        with Image.open(BytesIO(contents)) as pimg:
            nd_arr = np.array(pimg)
            return self.parse_img_nd_to_text(nd_arr, mode)

    def parse_file_to_text(self, filename):
        return self.parse_img_to_text(filename, DEFAULT_TBPU)



    def parse_img_file(self, img_path):
        img = Image.open(img_path)
        return self.parse_img_nd(np.array(img))

    def parse_img_nd(self, nd_arr):
        if nd_arr is None:
            return []
        result = paddle.ocr(nd_arr, cls=False)
        resp = []
        for idx in range(len(result)):
            res = result[idx]
            if res is not None:
                for line in res:
                    text, score = line[1]
                    resp.append({
                        'box': line[0],
                        'text': text,
                        'score': score
                    })
        return resp

    def parse_img_nd_to_text(self, nd_arr, tbpu):
        resp = self.parse_img_nd(nd_arr)
        # 计算平均置信度
        score, num = 0, 0
        for r in resp:
            score += r["score"]
            num += 1
        if num > 0:
            score /= num

        # 执行 tbpu
        paser = get_parser(tbpu)
        if paser is not None and resp is not None and len(resp) > 0:
            resp = paser.run(resp)
        resp = get_data_text(resp)
        return {
            'score': score,
            'data': resp
        }


    def parse_img_to_text(self, img_path, tbpu):
        resp = self.parse_img_file(img_path)
        # 计算平均置信度
        score, num = 0, 0
        for r in resp:
            score += r["score"]
            num += 1
        if num > 0:
            score /= num

        # 执行 tbpu
        paser = get_parser(tbpu)
        if paser is not None and resp is not None and len(resp) > 0:
            resp = paser.run(resp)
        resp = get_data_text(resp)
        return {
            'score': score,
            'data': resp
        }
